#include "Broomie.h"
#include <broomie/libbrm.hpp>
#include <broomie/brmalgorithm.hpp>
#include <protocol/TBinaryProtocol.h>
#include <server/TThreadPoolServer.h>
#include <concurrency/ThreadManager.h>
#include <concurrency/PosixThreadFactory.h>
#include <transport/TServerSocket.h>
#include <transport/TBufferTransports.h>

using namespace apache::thrift;
using namespace apache::thrift::protocol;
using namespace apache::thrift::transport;
using namespace apache::thrift::server;
using namespace apache::thrift::concurrency;

using boost::shared_ptr;

class BroomieHandler : virtual public BroomieIf {
 public:
  BroomieHandler() {
    modelFactory = new broomie::ModelFactoryImpl();
    basePath = "/var/mixi/test/";
    cl = new broomie::Classifier(modelFactory, basePath);
  }

  void addTraining(const std::string& train) {
    std::vector<std::string> features;
    features = strSplit(train, "\t");
    std::string className = features.front();
    if(classNamesMap->find(className) == classNamesMap->end()){
      std::cerr << "unknown class:[" << className << "]" << std::endl;
      std::cerr << "check class names of test data or traning data"
                << std::endl;
      if(!cl->endMakingModel()){
        std::string errMessage = cl->traceErr();
        std::cerr << errMessage << std::endl;
      }
      delete modelFactory;
    }
    features.erase(features.begin());
    if(features.size() < 2) return;
    unsigned int featuresSiz = features.size();
    double docSiz = ( featuresSiz - 1) / 2;
    ++docSiz;
    broomie::Document *doc = new broomie::Document(static_cast<int>(docSiz));
    std::string feature;
    for(unsigned int i = 0; i < featuresSiz; i++){
      if(i % 2 == 0){
        feature = features[i];
      } else {
        double point = atof(features[i].c_str());
        doc->addFeature(feature, point);
      }
    }
    if(!cl->addTrainingData(className, *doc)){
      std::string errMessage = cl->traceErr();
      std::cout << errMessage << std::endl;
    }
    delete doc;
    fflush(stdout);
  }

  void beginMakingModel(const std::string& classNamesBuf) {
    std::vector<std::string> classNames;
    classNames.push_back("it");
    classNames.push_back("economic");
    classNames.push_back("sport");
    int classifierMethod = broomie::BAYES;
    cl->beginMakingModel(classNames, classifierMethod);
    std::string confPath = basePath + broomie::CONFIG_NAME;
    std::ofstream ofs(confPath.c_str());
    ofs.write(broomie::DEFINE_METHOD_NAME.c_str(),
              std::strlen(broomie::DEFINE_METHOD_NAME.c_str()));
    ofs.write("\t", 1);
    std::string methodBuf;
    if(classifierMethod == broomie::BAYES){
      methodBuf = broomie::METHOD_BAYES;
    } else if(classifierMethod == broomie::OLL){
      methodBuf = broomie::METHOD_OLL;
    } else if(classifierMethod == broomie::TINYSVM){
      methodBuf = broomie::METHOD_TINYSVM;
    } else {
      //return false;
    }
    ofs.write(methodBuf.c_str(), std::strlen(methodBuf.c_str()));
    ofs.write("\n", 1);
    ofs.write(broomie::DEFINE_BASE_DIR_NAME.c_str(),
              std::strlen(broomie::DEFINE_BASE_DIR_NAME.c_str()));
    ofs.write("\t", 1);
    ofs.write(".", std::strlen("."));
    ofs.write("\n", 1);
    ofs.close();
    classNamesMap = new std::map<std::string, int>();
    for(unsigned int i = 0; i < classNames.size(); i++){
      classNamesMap->insert(std::make_pair(classNames[i], 1));
    }
  }

  void endMakingModel() {
    if(!cl->endMakingModel()){
      std::string errMessage = cl->traceErr();
      std::cerr << errMessage << std::endl;
    }
  }

  void beginClassification() {
    std::string confPath = basePath + broomie::CONFIG_NAME;
    std::ifstream confifs(confPath.c_str(), std::ios::in);
    if(!confifs){
      std::cerr << "file open error: " << confPath << std::endl;
      return;
    }
    int classifierMethod = 0;
    std::string line;
    while(std::getline(confifs, line)){
      std::vector<std::string> features = strSplit(line, "\t");
      if(features[0] == broomie::DEFINE_METHOD_NAME){
        if(features[1] == broomie::METHOD_BAYES){
          classifierMethod = broomie::BAYES;
        } else if(features[1] == broomie::METHOD_OLL){
          classifierMethod = broomie::OLL;
        } else if(features[1] == broomie::METHOD_TINYSVM){
            classifierMethod = broomie::TINYSVM;
        } else {
          std::cout << features[1] << "\t" << broomie::METHOD_TINYSVM << std::endl;
          std::cerr << "chose classifier method error: "
                    << features[1] <<  std::endl;
          return;
        }
      }
    }
    if(!cl->beginClassification(classifierMethod)){
        std::string errMessage = cl->traceErr();
        std::cerr << errMessage << std::endl;
        return;
    }
    confifs.close();
  }

  void classify(std::string& _return, const std::string& input) {
    std::vector<std::string> features;
    features = strSplit(input, "\t");
    if(features.size() < 2) return;
    std::string docId = features.front();
    features.erase(features.begin());

    unsigned int featuresSiz = features.size();
    double docSiz = ( featuresSiz - 1) / 2;
    broomie::Document *doc =
      new broomie::Document(static_cast<int>(docSiz));
    std::string feature("");
    for(unsigned int i = 0; i < featuresSiz - 1; i++){
      if(i % 2 == 0){
        feature = features[i];
      } else {
        double point = atof(features[i].c_str());
        doc->addFeature(feature, point);
      }
    }
    broomie::ResultSet *rs = cl->classify(*doc);
    for(int i = 0; i < rs->getResultSetNum(); i++){
      float point;
      std::string resultClassName = rs->getResult(i, point);
      char pointStr[50];
      sprintf(pointStr, "%f", point);
      _return.append(resultClassName);
      _return.append("\t");
      _return.append(pointStr);
      _return.append("\n");
    }
  }

  void endClassification() {
    if(!cl->endClassification()){
      std::string errMessage = cl->traceErr();
      std::cerr << errMessage << std::endl;
    }
  }

  private:
  broomie::ModelFactoryImpl* modelFactory;
  broomie::Classifier* cl;
  std::map<std::string, int>* classNamesMap;
  std::string basePath;

  std::vector<std::string> strSplit(std::string str, std::string delim)
  {
    std::vector<std::string> result;
    uint64_t cutAt;
    while((cutAt = str.find_first_of(delim)) != str.npos){
      if(cutAt > 0) result.push_back(str.substr(0, cutAt));
      str = str.substr(cutAt + 1);
    }
    if(str.length() > 0) result.push_back(str);
    return result;
  }
};

int main(int argc, char **argv) {
  int port = 9090;
  int workerNum = 50;
  shared_ptr<BroomieHandler> handler(new BroomieHandler());
  shared_ptr<TProcessor> processor(new BroomieProcessor(handler));
  shared_ptr<TServerTransport> serverTransport(new TServerSocket(port));
  shared_ptr<TTransportFactory> transportFactory(new TBufferedTransportFactory());
  shared_ptr<TProtocolFactory> protocolFactory(new TBinaryProtocolFactory());

  shared_ptr<ThreadManager> threadManager =
    ThreadManager::newSimpleThreadManager(workerNum);
  shared_ptr<PosixThreadFactory> threadFactory =
    shared_ptr<PosixThreadFactory>(new PosixThreadFactory());
  threadManager->threadFactory(threadFactory);
  threadManager->start();
  TThreadPoolServer server(processor, serverTransport, transportFactory,
                           protocolFactory, threadManager);
  server.serve();
  return 0;
}

